#include <tinx/memory.h>
#include <tinx/task.h>
#include <tinx/interrupt.h>
#include <tinx/debug.h>
#include <tinx/fs.h>
#include <tinx/arena.h>
#include <tinx/string.h>
#include <tinx/stdlib.h>

enum
{
    AT_NULL = 0,    /* end of vector */
    AT_IGNORE,      /* entry should be ignored */
    AT_EXECFD,      /* file descriptor of program */
    AT_PHDR,        /* program headers for program */
    AT_PHENT,       /* size of program header entry */
    AT_PHNUM,       /* number of program headers */
    AT_PAGESZ,      /* system page size */
    AT_BASE,        /* base address of interpreter */
    AT_FLAGS,       /* flags */
    AT_ENTRY,       /* entry point of program */
    AT_NOTELF,      /* program is not ELF */
    AT_UID,         /* real uid */
    AT_EUID,        /* effective uid */
    AT_GID,         /* real gid */
    AT_EGID,        /* effective gid */
    AT_PLATFORM,
    AT_HWCAP,
    AT_EXECFN = 31
};

typedef u32 Elf32_Word;
typedef u32 Elf32_Addr;
typedef u32 Elf32_Off;
typedef u16 Elf32_Half;

typedef struct Elf32_Auxv
{
    u32 type;
    u32 value;
} _packed Elf32_Auxv;

// ELF 文件标记
typedef struct ELFIdent
{
    u8 ei_magic[4];    // 内容为 0x7F, E, L, F
    u8 ei_class;       // 文件种类 1-32位，2-64 位
    u8 ei_data;        // 标记大小端 1-小端，2-大端
    u8 ei_version;     // 与 e_version 一样，必须为 1
    u8 ei_pad[16 - 7]; // 占满 16 个字节
} _packed ELFIdent;

// ELF 文件头
typedef struct Elf32_Ehdr
{
    ELFIdent e_ident;       // ELF 文件标记，文件最开始的 16 个字节
    Elf32_Half e_type;      // 文件类型，见 Etype
    Elf32_Half e_machine;   // 处理器架构类型，标记运行的 CPU，见 EMachine
    Elf32_Word e_version;   // 文件版本，见 EVersion
    Elf32_Addr e_entry;     // 程序入口地址
    Elf32_Off e_phoff;      // program header offset 程序头表在文件中的偏移量
    Elf32_Off e_shoff;      // section header offset 节头表在文件中的偏移量
    Elf32_Word e_flags;     // 处理器特殊标记
    Elf32_Half e_ehsize;    // ELF header size ELF 文件头大小
    Elf32_Half e_phentsize; // program header entry size 程序头表入口大小
    Elf32_Half e_phnum;     // program header number 程序头数量
    Elf32_Half e_shentsize; // section header entry size 节头表入口大小
    Elf32_Half e_shnum;     // section header number 节头表数量
    Elf32_Half e_shstrndx;  // 节字符串表在节头表中的索引
} Elf32_Ehdr;

// ELF 文件类型
enum Etype
{
    ET_NONE = 0,        // 无类型
    ET_REL = 1,         // 可重定位文件
    ET_EXEC = 2,        // 可执行文件
    ET_DYN = 3,         // 动态链接库
    ET_CORE = 4,        // core 文件，未说明，占位
    ET_LOPROC = 0xff00, // 处理器相关低值
    ET_HIPROC = 0xffff, // 处理器相关高值
};

// ELF 机器(CPU)类型
enum EMachine
{
    EM_NONE = 0,  // No machine
    EM_M32 = 1,   // AT&T WE 32100
    EM_SPARC = 2, // SPARC
    EM_386 = 3,   // Intel 80386
    EM_68K = 4,   // Motorola 68000
    EM_88K = 5,   // Motorola 88000
    EM_860 = 7,   // Intel 80860
    EM_MIPS = 8,  // MIPS RS3000
};

// ELF 文件版本
enum EVersion
{
    EV_NONE = 0,    // 不可用版本
    EV_CURRENT = 1, // 当前版本
};

// 程序头
typedef struct Elf32_Phdr
{
    Elf32_Word p_type;   // 段类型，见 SegmentType
    Elf32_Off p_offset;  // 段在文件中的偏移量
    Elf32_Addr p_vaddr;  // 加载到内存中的虚拟地址
    Elf32_Addr p_paddr;  // 加载到内存中的物理地址
    Elf32_Word p_filesz; // 文件中占用的字节数
    Elf32_Word p_memsz;  // 内存中占用的字节数
    Elf32_Word p_flags;  // 段标记，见 SegmentFlag
    Elf32_Word p_align;  // 段对齐约束
} Elf32_Phdr;

// 段类型
enum SegmentType
{
    PT_NULL = 0,    // 未使用
    PT_LOAD = 1,    // 可加载程序段
    PT_DYNAMIC = 2, // 动态加载信息
    PT_INTERP = 3,  // 动态加载器名称
    PT_NOTE = 4,    // 一些辅助信息
    PT_SHLIB = 5,   // 保留
    PT_PHDR = 6,    // 程序头表
    PT_LOPROC = 0x70000000,
    PT_HIPROC = 0x7fffffff,
};

enum SegmentFlag
{
    PF_X = 0x1, // 可执行
    PF_W = 0x2, // 可写
    PF_R = 0x4, // 可读
};

typedef struct Elf32_Shdr
{
    Elf32_Word sh_name;      // 节名，在字符串表中的索引
    Elf32_Word sh_type;      // 节类型，见 SectionType
    Elf32_Word sh_flags;     // 节标记，见 SectionFlag
    Elf32_Addr sh_addr;      // 节地址
    Elf32_Off sh_offset;     // 节在文件中的偏移量
    Elf32_Word sh_size;      // 节大小
    Elf32_Word sh_link;      // 保存了头表索引链接，与节类型相关
    Elf32_Word sh_info;      // 额外信息，与节类型相关
    Elf32_Word sh_addralign; // 地址对齐约束
    Elf32_Word sh_entsize;   // 子项入口大小
} Elf32_Shdr;

enum SectionType
{
    SHT_NULL = 0,            // 不可用
    SHT_PROGBITS = 1,        // 程序信息
    SHT_SYMTAB = 2,          // 符号表
    SHT_STRTAB = 3,          // 字符串表
    SHT_RELA = 4,            // 有附加重定位
    SHT_HASH = 5,            // 符号哈希表
    SHT_DYNAMIC = 6,         // 动态链接信息
    SHT_NOTE = 7,            // 标记文件信息
    SHT_NOBITS = 8,          // 该节文件中无内容
    SHT_REL = 9,             // 无附加重定位
    SHT_SHLIB = 10,          // 保留，用于非特定的语义
    SHT_DYNSYM = 11,         // 符号表
    SHT_LOPROC = 0x70000000, // 以下与处理器相关
    SHT_HIPROC = 0x7fffffff,
    SHT_LOUSER = 0x80000000,
    SHT_HIUSER = 0xffffffff,
};

enum SectionFlag
{
    SHF_WRITE = 0x1,           // 执行时可写
    SHF_ALLOC = 0x2,           // 执行时占用内存，有些节执行时可以不在内存中
    SHF_EXECINSTR = 0x4,       // 包含可执行的机器指令，节里有代码
    SHF_MASKPROC = 0xf0000000, // 保留，与处理器相关
};

typedef struct Elf32_Sym
{
    Elf32_Word st_name;  // 符号名称，在字符串表中的索引
    Elf32_Addr st_value; // 符号值，与具体符号相关
    Elf32_Word st_size;  // 符号的大小
    u8 st_info;          // 指定符号类型和约束属性，见 SymbolBinding
    u8 st_other;         // 为 0，无意义
    Elf32_Half st_shndx; // 符号对应的节表索引
} Elf32_Sym;

// 通过 info 获取约束
#define ELF32_ST_BIND(i) ((i) >> 4)
// 通过 info 获取类型
#define ELF32_ST_TYPE(i) ((i)&0xF)
// 通过 约束 b 和 类型 t 获取 info
#define ELF32_ST_INFO(b, t) (((b) << 4) + ((t)&0xf))

static bool elf_validate(Elf32_Ehdr *ehdr)
{
    // 不是 ELF 文件
    if (memcmp(&ehdr->e_ident, "\177ELF\1\1\1", 7))
        return false;

    // 不是可执行文件
    if (ehdr->e_type != ET_EXEC && ehdr->e_type != ET_DYN)
        return false;

    // 不是 386 平台文件
    if (ehdr->e_machine != EM_386)
        return false;

    // 版本不可识别
    if (ehdr->e_version != EV_CURRENT)
        return false;

    if (ehdr->e_phentsize != sizeof(Elf32_Phdr))
        return false;

    return true;
}

int count_argv(char *argv[])
{
    if (!argv)
        return 0;
    int nr = 0;
    while (argv[nr])
        nr++;
    return nr;
}

u32 new_auxv_entry(u32 top, u32 type, u32 value)
{
    top -= sizeof(Elf32_Auxv);
    Elf32_Auxv *auxv = (Elf32_Auxv *)top;
    auxv->type = type;
    auxv->value = value;
}

u32 copy_auxvs(u32 top, Elf32_Ehdr *ehdr, u32 base, char *execfn, u32 entry)
{
    task_t *task = running_task();
    u32 newtop = top;

    newtop = new_auxv_entry(newtop, AT_NULL, 0);
    newtop = new_auxv_entry(newtop, AT_UID, task->uid);
    newtop = new_auxv_entry(newtop, AT_EUID, 0);
    newtop = new_auxv_entry(newtop, AT_GID, 0);
    newtop = new_auxv_entry(newtop, AT_EGID, 0);
    newtop = new_auxv_entry(newtop, AT_PHDR, task->exec + ehdr->e_phoff);
    newtop = new_auxv_entry(newtop, AT_PHENT, sizeof(Elf32_Phdr));
    newtop = new_auxv_entry(newtop, AT_PHNUM, ehdr->e_phnum);
    newtop = new_auxv_entry(newtop, AT_PAGESZ, PAGE_SIZE);
    newtop = new_auxv_entry(newtop, AT_BASE, base);
    newtop = new_auxv_entry(newtop, AT_FLAGS, 0);
    newtop = new_auxv_entry(newtop, AT_ENTRY, entry);
    newtop = new_auxv_entry(newtop, AT_HWCAP, 0);
    newtop = new_auxv_entry(newtop, AT_EXECFN, (u32)execfn);

    return top - newtop;
}

static u32 copy_argv_envp_auxv(char *argv[], char *envp[], Elf32_Ehdr *ehdr, u32 base, u32 entry)
{
    task_t *task = running_task();

    // 用户栈
    char *utop = (char *)USER_STACK_TOP;
    // 缓冲区
    char *buf = (char *)alloc_kpage(4);
    char *buf_end = buf + PAGE_SIZE * 4;
    char *ktop = buf_end;

    u32 argc = count_argv(argv);
    u32 envc = count_argv(envp);

    char **argvk = (char **)alloc_kpage(1);
    char **envpk = argvk + argc + 1;

    int len = strlen(task->name) + 1;
    ktop -= len;
    utop -= len;
    memcpy(ktop, task->name, len);
    char *execfn = utop;

    // 拷贝 envs
    for (size_t i = 0; i < envc; i++)
    {
        char *env = envp[i];
        int len = strlen(env);
        ktop -= ++len;
        utop -= len;
        memcpy(ktop, env, len);
        envpk[i] = utop;
    }

    // 拷贝 args
    for (size_t i = 0; i < argc; i++)
    {
        char *arg = argv[i];
        int len = strlen(arg);
        ktop -= ++len;
        utop -= len;
        memcpy(ktop, arg, len);
        argvk[i] = utop;
    }

    int size = copy_auxvs((u32)ktop, ehdr, base, execfn, entry);
    ktop -= size;
    utop -= size;

    // 拷贝 envp
    int count = (envc + 1) * sizeof(char *);
    ktop -= count;
    utop -= count;
    memcpy(ktop, (void *)envpk, count);

    // 拷贝 argv
    count = (argc + 1) * sizeof(char *);
    ktop -= count;
    utop -= count;
    memcpy(ktop, (void *)argvk, count);

    // argc
    ktop -= sizeof(u32);
    utop -= sizeof(u32);
    *(u32 *)ktop = argc;

    // 拷贝栈
    memcpy(utop, ktop, buf_end - ktop);

    // 释放缓冲
    free_kpage((u32)buf, 4);
    free_kpage((u32)argvk, 1);

    return (u32)utop;
}

static void load_segment(inode_t *inode, Elf32_Phdr *phdr, u32 base)
{
    task_t *task = running_task();

    // 获得段加载在内存中的地址
    u32 vaddr = phdr->p_vaddr;
    vaddr += base;

    u32 page = vaddr & ~0xfff;

    u32 count = div_round_up(MAX(phdr->p_filesz, phdr->p_memsz) + (vaddr & 0xfff), PAGE_SIZE);

    // 映射页
    for (size_t i = 0; i < count; i++)
    {
        u32 addr = page + i * PAGE_SIZE;
        link_page(addr);
        if (phdr->p_flags & PF_W)
            continue;
        // 如果段不可写，则置为只读
        page_entry_t *entry = get_pte(addr, false);
        entry->write = false;
        flush_tlb(addr);
    }

    // 加载程序段
    inode->op->read(inode, (char *)vaddr, phdr->p_filesz, phdr->p_offset);
    // 初始化 .bss 段
    if (phdr->p_memsz > phdr->p_filesz)
    {
        u32 addr = vaddr + phdr->p_filesz;
        memset((char *)addr, 0, phdr->p_memsz - phdr->p_filesz);
    }

    if (phdr->p_flags == (PF_X | PF_R))
    {
        task->text = vaddr;
    }
    else if (phdr->p_flags == (PF_W | PF_R))
    {
        task->data = vaddr;
    }

    task->end = MAX(task->end, page + PAGE_SIZE * count);
}

static int load_interp_segment(inode_t *einode, Elf32_Phdr *ephdr, u32 *entry)
{
    task_t *task = running_task();

    char *file = NULL;
    int ret = EOK;

    file = kmalloc(ephdr->p_filesz);
    einode->op->read(einode, file, ephdr->p_filesz, ephdr->p_offset);
    inode_t *inode = namei(file);

    // 不存在
    if (!inode)
    {
        ret = -ENOEXEC;
        goto rollback;
    }

    // 不是文件
    if (!ISFILE(inode->mode))
    {
        ret = -ENOEXEC;
        goto rollback;
    }

    // 无权限
    if (!inode->op->permission(inode, P_EXEC))
    {
        ret = -ENOEXEC;
        goto rollback;
    }

    u32 addr = alloc_kpage(1);
    Elf32_Ehdr *ehdr = (Elf32_Ehdr *)addr;
    inode->op->read(inode, (char *)ehdr, sizeof(Elf32_Ehdr), 0);

    if (!elf_validate(ehdr))
    {
        free_kpage(addr, 1);
        ret = -ENOEXEC;
        goto rollback;
    }

    // 清理堆内存
    task->end = USER_EXEC_ADDR;
    sys_brk(task->end);

    Elf32_Phdr *phdr = (Elf32_Phdr *)(ehdr + 1);
    inode->op->read(inode, phdr, ehdr->e_phnum * ehdr->e_phentsize, ehdr->e_phoff);

    // 逐个检查
    Elf32_Phdr *ptr = phdr;
    for (size_t i = 0; i < ehdr->e_phnum; i++, ptr++)
    {
        if (ptr->p_type != PT_LOAD)
            continue;
        load_segment(inode, ptr, USER_INTERP_ADDR);
    }

    *entry = USER_INTERP_ADDR + ehdr->e_entry;

    free_kpage(addr, 1);

rollback:
    iput(inode);
    kfree(file);
    return ret;
}

int sys_execve(char *filename, char *argv[], char *envp[])
{
    task_t *task = running_task();
    inode_t *inode = namei(filename);
    u32 base = 0;
    u32 entry = 0;

    int ret;

    // 不存在
    if (!inode)
    {
        ret = -ENOEXEC;
        goto rollback;
    }

    // 不是文件
    if (!ISFILE(inode->mode))
    {
        ret = -ENOEXEC;
        goto rollback;
    }

    // 无权限
    if (!inode->op->permission(inode, P_EXEC))
    {
        ret = -ENOEXEC;
        goto rollback;
    }

    u32 addr = alloc_kpage(1);
    Elf32_Ehdr *ehdr = (Elf32_Ehdr *)addr;
    inode->op->read(inode, (char *)ehdr, sizeof(Elf32_Ehdr), 0);

    if (!elf_validate(ehdr))
    {
        free_kpage(addr, 1);
        ret = -ENOEXEC;
        goto rollback;
    }

    if (ehdr->e_type == ET_DYN)
        base = USER_EXEC_ADDR;

    // 加载所有段头，并且设置加载地址,为下面创建辅助向量做准备
    Elf32_Phdr *phdr = (Elf32_Phdr *)(ehdr + 1);
    inode->op->read(inode, phdr, ehdr->e_phnum * ehdr->e_phentsize, ehdr->e_phoff);

    u32 exec = USER_MEMORY_END;

    // 逐个检查
    Elf32_Phdr *ptr = phdr;
    for (size_t i = 0; i < ehdr->e_phnum; i++, ptr++)
    {
        if (ptr->p_type == PT_PHDR || ptr->p_type == PT_LOAD)
            // 设置新的执行地址，并不是入口地址
            exec = MIN(exec, ptr->p_vaddr + base);
        if (ptr->p_type != PT_INTERP)
            continue;
        ret = load_interp_segment(inode, ptr, &entry);
        if (ret < 0)
            goto rollback;
    }

    if (!entry)
    {
        // 清理堆内存
        task->end = USER_EXEC_ADDR;
        sys_brk(task->end);
    }

    task->exec = exec;

    strncpy(task->name, filename, 16);

    // 拷贝 argv envp 和 辅助向量
    u32 top = copy_argv_envp_auxv(argv, envp, ehdr, entry ? USER_INTERP_ADDR : base, base + ehdr->e_entry);

    // 开始处理所有段
    ptr = (Elf32_Phdr *)(ehdr + 1);
    for (size_t i = 0; i < ehdr->e_phnum; i++, ptr++)
    {
        if (ptr->p_type == PT_LOAD || ptr->p_type == PT_PHDR)
            load_segment(inode, ptr, base);
    }

    // 设置新的执行文件 inode
    iput(task->iexec);
    task->iexec = inode;

    // 设置新的 brk
    sys_brk(task->end);

    if (!entry)
        entry = base + ehdr->e_entry;

    free_kpage(addr, 1);

    for (int i = 0; i < OPEN_NR; i++)
    {
        file_t *file = task->files[i];
        if (!file)
            continue;
        if (file->flags & O_CLOEXEC)
            put_file(file);
    }

    // 直接返回到程序
    addr = (u32)task + PAGE_SIZE - sizeof(intr_frame_t);
    intr_frame_t *iframe = (intr_frame_t *)addr;
    iframe->eip = entry;
    iframe->esp = top;

    DEBUGK("%s entry 0x%p\n", task->name, ehdr->e_entry);

    asm volatile(
        "movl %%eax, %%esp\n"
        "jmp interrupt_exit\n" :: "a"(iframe));

    return EOK;

rollback:
    iput(inode);
    return ret;
}